package algorithms.weekly.class01;

import java.util.ArrayList;
import java.util.List;

/**
 * @author guichang
 * @date 2021/12/26
 */

@SuppressWarnings("all")
public class Code1_保留黑节点二叉树 {

    private static class Node {
        int value; // 值
        boolean retain;// 是否保留 true-保留
        List<Node> nodes; // 孩子节点列表

        public Node(int value, boolean retain, List<Node> nodes) {
            this.value = value;
            this.retain = retain;
            this.nodes = nodes;
        }
    }

    public static Node retain(Node head) {
        if (head == null) {
            return null;
        }
        return process(head);
    }

    /**
     * 当前节点x调整完后，返回新头部
     */
    public static Node process(Node x) {
        // 当前节点无孩子节点，当前节点保留就返回当前节点，否则返回null
        if (x.nodes == null || x.nodes.isEmpty()) {
            return x.retain ? x : null;
        }
        // 遍历孩子节点，保留非空的
        List<Node> newNodes = new ArrayList<>();
        for (Node node : x.nodes) {
            Node n = process(node);
            if (n != null) {
                newNodes.add(n);
            }
        }
        // 新孩子节点不为空或者自己要保留则替换新孩子节点
        if (!newNodes.isEmpty() || x.retain) {
            x.nodes = newNodes;
            return x;
        }
        // 否则就是什么都不保留了
        return null;
    }


}